import torch
from varflow.distributions import Distribution, ConditionalBernoulli
from varflow.flows import FlowLayer

class Rejection(FlowLayer):
    """A deep acceptance rejection sampling flow."""

    def __init__(self, sampler: ConditionalBernoulli, verbose=False, leak=0.01, rev=1):
        super(Rejection, self).__init__()
        self.sampler = sampler
        self.verbose = verbose
        self.leak = leak
        self.rev = rev
        self.register_buffer('prior', torch.tensor(0.5))

    def log_prob(self, x):
        posterior = self.sampler.probs(context=x)[..., 0]
        G = self.base_dist.sample(x.shape[0])
        G = (1 + self.rev) * G.detach() - self.rev * G
        prior = self.sampler.probs(context=G).mean()
        self.prior = prior.detach()
        log_prior = torch.log(prior)
        return self.base_dist.log_prob(x) + posterior.log() - log_prior

    def sample(self, num_samples):
        z = []
        if self.verbose:
            from tqdm import tqdm
            pbar = tqdm(total=int(num_samples))
        while len(z) < num_samples:
            x = self.base_dist.sample(num_samples)
            k = self.sampler.sample(context=x)[..., 0].bool()
            z = x[k] if isinstance(z, list) else torch.cat([z, x[k]])
            if self.verbose:
                pbar.update(k.sum().item())
        return z[:num_samples]
